在昨天的學習中,我們深入了解了 Locust 的 HTTPClient 功能,學會了如何發送各種請求和處理回應。今天,我們將進一步探討在實際測試中非常重要的主題:Cookie 處理和 Session 管理。透過掌握這些技巧,您將能夠模擬真實的使用者會話行為,測試需要認證的應用程式,以及處理複雜的狀態管理場景。
在深入 Locust 的實作之前,讓我們先回顧一下 Session 和 Cookie 的基本概念:
Cookie 是伺服器傳送到使用者瀏覽器並儲存在本地的小型資料片段。在後續的請求中,瀏覽器會自動將 Cookie 傳送回伺服器,讓伺服器能夠識別使用者身份或維持狀態。
Session 是伺服器端用來儲存使用者狀態的機制。通常,伺服器會產生一個唯一的 Session ID,並透過 Cookie 傳送給客戶端。客戶端在後續請求中帶上這個 Session ID,伺服器就能識別並恢復使用者的狀態。
在壓力測試中,正確處理 Cookie 和 Session 至關重要,因為:
Locust 的 HTTPClient 基於 Python 的 requests
函式庫,它會自動處理 Cookie,就像瀏覽器一樣。這意味著:
Set-Cookie
標頭時,Cookie 會自動儲存讓我們看一個基本的例子:
在下面的程式碼中,我們利用 Locust 的自動 Cookie 管理機制來處理使用者登入和認證流程。透過 on_start()
生命週期方法執行一次性的登入操作,然後在 @task
裝飾的方法中模擬需要認證的操作。Locust 的 HTTPClient 會自動處理伺服器設定的 Session Cookie,讓後續的請求能夠維持登入狀態,這完全模擬了瀏覽器的行為:
from locust import HttpUser, task, between
class AutoCookieUser(HttpUser):
wait_time = between(1, 3)
host = "http://localhost:8080"
def on_start(self):
"""使用者開始時執行登入"""
# 發送登入請求
response = self.client.post("/login", json={
"username": "test_user",
"password": "password123"
})
# 伺服器會在回應中設定 session cookie
# Locust 會自動儲存這個 cookie
if response.status_code == 200:
print("登入成功,Session Cookie 已自動儲存")
@task
def access_protected_resource(self):
"""訪問需要認證的資源"""
# Cookie 會自動包含在請求中
response = self.client.get("/profile")
if response.status_code == 200:
print("成功訪問受保護的資源")
elif response.status_code == 401:
print("認證失敗")
雖然自動處理很方便,但有時我們需要更精細的控制。Locust 提供了多種方式來手動管理 Cookie。
class CookieInspectionUser(HttpUser):
wait_time = between(1, 2)
@task
def inspect_cookies(self):
"""檢查當前的 cookies"""
# 獲取所有 cookies
all_cookies = self.client.cookies
print(f"當前 Cookies: {dict(all_cookies)}")
# 檢查特定 cookie 是否存在
if 'session_token' in self.client.cookies:
token_value = self.client.cookies['session_token']
print(f"Session Token: {token_value}")
# 遍歷所有 cookies
for cookie_name, cookie_value in self.client.cookies.items():
print(f"{cookie_name}: {cookie_value}")
以下程式碼展示了如何完全控制 Cookie 的管理。我們使用 self.client.cookies.clear()
清除所有現有的 Cookie,然後透過 self.client.cookies.set()
方法手動設定各種 Cookie。這種方式特別適用於需要模擬特定使用者偏好設定、測試不同地區設定、或需要精確控制 Cookie 內容的測試場景:
class ManualCookieUser(HttpUser):
wait_time = between(1, 2)
def on_start(self):
"""手動設定 cookies"""
# 清除所有現有 cookies
self.client.cookies.clear()
# 設定單個 cookie
self.client.cookies.set('user_preference', 'dark_mode')
# 設定多個 cookies
cookies_to_set = {
'language': 'zh-TW',
'timezone': 'Asia/Taipei',
'session_id': 'custom_session_123'
}
for name, value in cookies_to_set.items():
self.client.cookies.set(name, value)
@task
def send_with_custom_cookies(self):
"""發送帶有自訂 cookies 的請求"""
# 方法 1: 使用已設定的 cookies(自動帶上)
response = self.client.get("/api/data")
# 方法 2: 為特定請求添加額外的 cookies
response = self.client.get("/api/special",
cookies={'temporary': 'temp_value'})
@task
def manipulate_cookies(self):
"""Cookie 的各種操作"""
# 更新 cookie 值
if 'theme' in self.client.cookies:
self.client.cookies['theme'] = 'light_mode'
# 刪除特定 cookie
if 'temp_data' in self.client.cookies:
del self.client.cookies['temp_data']
# 複製 cookies 用於特殊請求
custom_cookies = dict(self.client.cookies)
custom_cookies['extra_param'] = 'special_value'
response = self.client.get("/api/endpoint", cookies=custom_cookies)
現在讓我們看看如何在 Locust 中實作完整的 Session 管理,包括登入、維持狀態和登出。
class SessionManagementUser(HttpUser):
wait_time = between(1, 3)
host = "http://localhost:8080"
def on_start(self):
"""初始化 session"""
self.session_active = False
self.login()
def login(self):
"""執行登入流程"""
print("執行登入...")
response = self.client.post("/login", json={
"username": "test_user",
"password": "password123"
})
if response.status_code == 200:
login_data = response.json()
self.session_token = login_data.get("session_token")
self.session_active = True
print(f"登入成功!Session Token: {self.session_token}")
# Cookie 會自動設定,但我們也可以手動檢查
if 'session_token' in self.client.cookies:
print("Session Cookie 已設定")
else:
print(f"登入失敗: {response.status_code}")
@task(5)
def perform_authenticated_action(self):
"""執行需要認證的操作"""
if not self.session_active:
self.login()
return
# 執行各種需要認證的操作
actions = [
("/profile", "GET"),
("/cart", "GET"),
("/orders", "GET")
]
endpoint, method = random.choice(actions)
if method == "GET":
response = self.client.get(endpoint)
if response.status_code == 401:
print("Session 已過期,重新登入")
self.session_active = False
self.login()
elif response.status_code == 200:
print(f"成功訪問 {endpoint}")
@task(1)
def logout(self):
"""登出並清理 session"""
if self.session_active:
response = self.client.post("/logout")
if response.status_code == 200:
print("登出成功")
self.session_active = False
# Cookie 會自動清除(如果伺服器正確設定)
# 重新登入以繼續測試
self.login()
def on_stop(self):
"""測試結束時清理"""
if self.session_active:
self.client.post("/logout")
print("測試結束,已登出")
在某些測試場景中,我們可能需要模擬一個使用者管理多個 session,或者在不同的 session 間切換。
class MultiSessionUser(HttpUser):
wait_time = between(2, 4)
host = "http://localhost:8080"
def on_start(self):
"""初始化多個 session"""
self.sessions = {}
self.current_session = None
# 建立多個不同的 session
for i in range(3):
self.create_session(f"user_{i}")
def create_session(self, username):
"""建立新的 session"""
# 暫存當前 cookies
original_cookies = dict(self.client.cookies)
# 清除 cookies 以建立新 session
self.client.cookies.clear()
# 登入
response = self.client.post("/login", json={
"username": "test_user",
"password": "password123"
})
if response.status_code == 200:
# 儲存這個 session 的 cookies
self.sessions[username] = dict(self.client.cookies)
print(f"Session 建立成功: {username}")
# 恢復原始 cookies
self.client.cookies.clear()
for name, value in original_cookies.items():
self.client.cookies.set(name, value)
def switch_session(self, username):
"""切換到指定的 session"""
if username in self.sessions:
# 清除當前 cookies
self.client.cookies.clear()
# 載入目標 session 的 cookies
for name, value in self.sessions[username].items():
self.client.cookies.set(name, value)
self.current_session = username
print(f"切換到 session: {username}")
@task
def random_session_action(self):
"""隨機切換 session 並執行操作"""
# 隨機選擇一個 session
username = random.choice(list(self.sessions.keys()))
self.switch_session(username)
# 執行操作
response = self.client.get("/profile")
if response.status_code == 200:
profile = response.json()
print(f"{username} 的個人資料: {profile}")
讓我們透過一個完整的購物車流程來展示 Session 管理的實際應用:
class ShoppingCartUser(HttpUser):
wait_time = between(1, 3)
host = "http://localhost:8080"
def on_start(self):
"""初始化購物流程"""
self.cart_items = 0
self.total_amount = 0
self.order_history = []
# 登入
self.login()
def login(self):
"""登入並建立 session"""
response = self.client.post("/login", json={
"username": "test_user",
"password": "password123"
})
if response.status_code == 200:
print("登入成功,開始購物")
# Session cookie 會自動設定
@task(10)
def browse_and_add_to_cart(self):
"""瀏覽商品並加入購物車"""
# 1. 瀏覽商品列表
products_response = self.client.get("/products")
if products_response.status_code == 200:
products = products_response.json()["products"]
# 2. 選擇一個商品
selected_product = random.choice(products)
# 3. 查看商品詳情
detail_response = self.client.get(f"/products/{selected_product['id']}")
if detail_response.status_code == 200:
# 4. 加入購物車
cart_item = {
"name": selected_product["name"],
"price": selected_product["price"],
"quantity": random.randint(1, 3)
}
add_response = self.client.post("/cart/add", json=cart_item)
if add_response.status_code == 200:
result = add_response.json()
self.cart_items = result["cart_items"]
# Session 會自動維持,購物車資料會保存
print(f"添加 {cart_item['name']} 到購物車,"
f"目前有 {self.cart_items} 項商品")
@task(5)
def view_cart(self):
"""查看購物車"""
# Session cookie 會自動帶上,所以能看到正確的購物車內容
response = self.client.get("/cart")
if response.status_code == 200:
cart = response.json()
self.total_amount = cart["total"]
print(f"購物車狀態: {cart['count']} 項商品,總價 ${cart['total']:.2f}")
@task(2)
def checkout(self):
"""結帳"""
if self.cart_items == 0:
print("購物車是空的,先添加商品")
return
# 結帳請求會使用當前 session 的購物車資料
checkout_response = self.client.post("/checkout")
if checkout_response.status_code == 200:
order = checkout_response.json()
self.order_history.append(order["order_id"])
print(f"訂單 {order['order_id']} 建立成功!總額: ${order['total']:.2f}")
# 重置購物車狀態
self.cart_items = 0
self.total_amount = 0
elif checkout_response.status_code == 400:
print("結帳失敗:購物車為空")
@task(1)
def view_order_history(self):
"""查看訂單歷史"""
if self.order_history:
# 查看最近的訂單
latest_order = self.order_history[-1]
response = self.client.get(f"/orders/{latest_order}")
if response.status_code == 200:
order = response.json()
print(f"訂單 {order['order_id']} 狀態: {order['status']}")
在大規模測試中,我們可能希望在多個虛擬使用者間共享或重用 session,以減少登入請求的數量:
class PersistentSessionUser(HttpUser):
wait_time = between(1, 3)
host = "http://localhost:8080"
# 類別層級的 session pool
session_pool = {}
pool_lock = threading.Lock()
def on_start(self):
"""從 pool 獲取或建立 session"""
self.get_or_create_session()
def get_or_create_session(self):
"""獲取可用的 session 或建立新的"""
with self.pool_lock:
# 嘗試獲取現有的 session
available_sessions = [
token for token, in_use in self.session_pool.items()
if not in_use
]
if available_sessions:
# 使用現有的 session
session_token = available_sessions[0]
self.session_pool[session_token] = True # 標記為使用中
# 設定 cookie
self.client.cookies.set('session_token', session_token)
print(f"重用現有 session: {session_token[:8]}...")
# 驗證 session 是否有效
if not self.verify_session():
self.create_new_session()
else:
# 建立新的 session
self.create_new_session()
def create_new_session(self):
"""建立新的 session 並加入 pool"""
response = self.client.post("/login", json={
"username": f"test_user_{random.randint(1, 100)}",
"password": "password123"
})
if response.status_code == 200:
session_data = response.json()
session_token = session_data["session_token"]
with self.pool_lock:
self.session_pool[session_token] = True # 標記為使用中
print(f"建立新 session: {session_token[:8]}...")
def verify_session(self):
"""驗證當前 session 是否有效"""
response = self.client.get("/profile")
return response.status_code == 200
@task
def perform_action(self):
"""執行需要認證的操作"""
response = self.client.get("/api/data")
if response.status_code == 401:
print("Session 失效,重新獲取")
self.get_or_create_session()
elif response.status_code == 200:
print("操作成功")
def on_stop(self):
"""釋放 session 回 pool"""
session_token = self.client.cookies.get('session_token')
if session_token:
with self.pool_lock:
if session_token in self.session_pool:
self.session_pool[session_token] = False # 標記為可用
print(f"釋放 session: {session_token[:8]}...")
class AdvancedCookieUser(HttpUser):
wait_time = between(1, 2)
@task
def handle_cookie_attributes(self):
"""處理 Cookie 的各種屬性"""
# 發送請求並檢查回應中的 Set-Cookie
response = self.client.get("/")
# 檢查 Set-Cookie 標頭
set_cookie_header = response.headers.get('Set-Cookie')
if set_cookie_header:
print(f"收到 Set-Cookie: {set_cookie_header}")
# 解析 cookie 屬性
if 'HttpOnly' in set_cookie_header:
print("Cookie 設定為 HttpOnly")
if 'Secure' in set_cookie_header:
print("Cookie 設定為 Secure(僅 HTTPS)")
if 'SameSite' in set_cookie_header:
print("Cookie 設定了 SameSite 屬性")
class CookieExpiryUser(HttpUser):
wait_time = between(1, 2)
def on_start(self):
"""初始化並設定 session 過期時間"""
self.session_created_at = time.time()
self.session_timeout = 300 # 5 分鐘
self.login()
def check_session_expiry(self):
"""檢查 session 是否過期"""
current_time = time.time()
elapsed = current_time - self.session_created_at
if elapsed > self.session_timeout:
print("Session 可能已過期,重新登入")
self.login()
self.session_created_at = current_time
return False
return True
@task
def authenticated_request(self):
"""執行需要認證的請求"""
# 先檢查 session 是否可能過期
if not self.check_session_expiry():
return
response = self.client.get("/api/protected")
if response.status_code == 401:
print("Session 確實已過期,重新登入")
self.login()
self.session_created_at = time.time()
以下是各個功能類別的執行命令,您可以根據需要選擇對應的測試類別:
locust -f demo.py --headless --users 3 --spawn-rate 1 --run-time 20s --host http://localhost:8080 --class-picker CookieSessionUser --only-summary
locust -f demo.py --headless --users 2 --spawn-rate 1 --run-time 15s --host http://localhost:8080 --class-picker ManualCookieUser --only-summary
locust -f demo.py --headless --users 1 --spawn-rate 1 --run-time 10s --host http://localhost:8080 --class-picker AdvancedCookieHandling --only-summary
class OptimizedSessionUser(HttpUser):
def on_start(self):
"""最佳化的 session 初始化"""
max_retries = 3
retry_count = 0
while retry_count < max_retries:
try:
response = self.client.post("/login", json={
"username": "test_user",
"password": "password123"
}, timeout=5)
if response.status_code == 200:
print("登入成功")
break
except Exception as e:
print(f"登入失敗 (嘗試 {retry_count + 1}/{max_retries}): {e}")
retry_count += 1
time.sleep(1)
@task
def secure_cookie_handling(self):
"""安全的 Cookie 處理"""
# 避免在日誌中顯示敏感的 cookie 值
sensitive_cookies = ['session_token', 'auth_token', 'csrf_token']
for cookie_name in self.client.cookies:
if cookie_name in sensitive_cookies:
print(f"{cookie_name}: [REDACTED]")
else:
print(f"{cookie_name}: {self.client.cookies[cookie_name]}")
class SessionPoolManager:
"""集中管理 session 池"""
_instance = None
_lock = threading.Lock()
def __new__(cls):
if cls._instance is None:
with cls._lock:
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._instance.sessions = {}
return cls._instance
def get_session(self):
"""獲取可用的 session"""
# 實作 session 分配邏輯
pass
def return_session(self, session_token):
"""歸還 session"""
# 實作 session 回收邏輯
pass
class RobustSessionUser(HttpUser):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.last_auth_check = time.time()
self.auth_check_interval = 60 # 每分鐘檢查一次
@task
def authenticated_action(self):
"""定期檢查並更新 session"""
current_time = time.time()
if current_time - self.last_auth_check > self.auth_check_interval:
# 發送心跳請求保持 session 活躍
self.client.get("/api/heartbeat")
self.last_auth_check = current_time
# 執行實際操作
response = self.client.get("/api/data")
class DistributedSessionUser(HttpUser):
"""適用於分散式測試的 session 管理"""
def on_start(self):
"""使用外部儲存同步 session"""
# 可以使用 Redis、資料庫等儲存 session
# 這裡用檔案系統作為簡單示例
session_file = f"/tmp/locust_sessions_{hash(self) % 100}.json"
try:
with open(session_file, 'r') as f:
session_data = json.load(f)
self.client.cookies.set('session_token', session_data['token'])
except FileNotFoundError:
# 建立新 session 並儲存
self.create_and_save_session(session_file)
class CSRFProtectedUser(HttpUser):
def on_start(self):
"""處理 CSRF 保護"""
# 先獲取 CSRF token
response = self.client.get("/")
# 從回應中提取 CSRF token(根據實際情況調整)
if 'X-CSRF-Token' in response.headers:
self.csrf_token = response.headers['X-CSRF-Token']
elif 'csrf_token' in response.cookies:
self.csrf_token = response.cookies['csrf_token']
# 在後續請求中包含 CSRF token
self.client.headers.update({'X-CSRF-Token': self.csrf_token})
class EfficientSessionUser(HttpUser):
# 類別變數,所有實例共享
_shared_session_token = None
_session_lock = threading.Lock()
def on_start(self):
"""優化登入流程"""
with self._session_lock:
if self._shared_session_token is None:
# 只有第一個使用者需要登入
response = self.client.post("/login", json={
"username": "shared_user",
"password": "password123"
})
if response.status_code == 200:
self._shared_session_token = response.json()["session_token"]
# 所有使用者使用相同的 session
if self._shared_session_token:
self.client.cookies.set('session_token', self._shared_session_token)
class SmartSessionUser(HttpUser):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.consecutive_failures = 0
self.max_failures = 3
@task
def smart_request(self):
"""智能處理 session 失效"""
response = self.client.get("/api/data")
if response.status_code == 401:
self.consecutive_failures += 1
if self.consecutive_failures >= self.max_failures:
# 多次失敗後才重新登入
print("多次認證失敗,重新登入")
self.login()
self.consecutive_failures = 0
else:
self.consecutive_failures = 0
今天我們深入學習了 Locust 中的 Cookie 處理和 Session 管理技術:
透過掌握這些技術,您現在能夠:
Cookie 和 Session 管理是壓力測試中不可或缺的技能,特別是在測試現代 Web 應用程式時。正確的 Session 管理不僅能讓測試更接近真實場景,還能提高測試的效率和準確性。